package com.computer.fundamentals.datastructure;

import com.computer.fundamentals.datastructure.node.BinaryTreeNode;
import com.computer.util.Constant;
import com.computer.util.UniversalMethod;

import java.util.*;

public class BinaryTree {

    public BinaryTreeNode root;

    public BinaryTree(BinaryTreeNode root_) {
        this.root = root_;
    }

    /**
     * 先序遍历（递归式）
     */
    public List<Integer> preorderByRecurrence(List<Integer> nums, BinaryTreeNode root) {
        if (root == null) {
            return nums;
        }
        nums.add(root.val);
        preorderByRecurrence(nums, root.left);
        preorderByRecurrence(nums, root.right);
        return nums;
    }

    /**
     * 中序遍历（递归式）
     */
    public List<Integer> inorderByRecurrence(List<Integer> nums, BinaryTreeNode root) {
        if (root == null) {
            return nums;
        }
        inorderByRecurrence(nums, root.left);
        nums.add(root.val);
        inorderByRecurrence(nums, root.right);
        return nums;
    }

    /**
     * 后序遍历（递归式）
     */
    public List<Integer> postorderByRecurrence(List<Integer> nums, BinaryTreeNode root) {
        if (root == null) {
            return nums;
        }
        postorderByRecurrence(nums, root.left);
        postorderByRecurrence(nums, root.right);
        nums.add(root.val);
        return nums;
    }

    /**
     * 颜色标记类，辅助迭代模拟递归
     */
    static class ColorFlag {
        BinaryTreeNode node;
        String color;

        public ColorFlag(BinaryTreeNode node, String color) {
            this.node = node;
            this.color = color;
        }
    }

    /**
     * 先序遍历（迭代式）
     */
    public List<Integer> preorderByIteration() {
        List<Integer> nums = new ArrayList<>();
        if (root == null) {
            return nums;
        }
        Stack<ColorFlag> stack = new Stack<>();
        stack.push(new ColorFlag(root, Constant.WHITE));
        while (!stack.isEmpty()) {
            ColorFlag colorFlag = stack.pop();
            BinaryTreeNode node = colorFlag.node;
            String color = colorFlag.color;
            if (color.equals(Constant.WHITE)) {
                if (node.right != null) {
                    stack.push(new ColorFlag(node.right, Constant.WHITE));
                }
                if (node.left != null) {
                    stack.push(new ColorFlag(node.left, Constant.WHITE));
                }
                stack.push(new ColorFlag(node, Constant.BLACK));
            }else {
                nums.add(node.val);
            }
        }
        return nums;
    }

    /**
     * 中序遍历（迭代式）
     */
    public List<Integer> inorderByIteration() {
        List<Integer> nums = new ArrayList<>();
        if (root == null) {
            return nums;
        }
        Stack<ColorFlag> stack = new Stack<>();
        stack.push(new ColorFlag(root, Constant.WHITE));
        while (!stack.isEmpty()) {
            ColorFlag colorFlag = stack.pop();
            BinaryTreeNode node = colorFlag.node;
            String color = colorFlag.color;
            if (color.equals(Constant.WHITE)) {
                if (node.right != null) {
                    stack.push(new ColorFlag(node.right, Constant.WHITE));
                }
                stack.push(new ColorFlag(node, Constant.BLACK));
                if (node.left != null) {
                    stack.push(new ColorFlag(node.left, Constant.WHITE));
                }
            }else {
                nums.add(node.val);
            }
        }
        return nums;
    }

    /**
     * 后序遍历（迭代式）
     */
    public List<Integer> postorderByIteration() {
        List<Integer> nums = new ArrayList<>();
        if (root == null) {
            return nums;
        }
        Stack<ColorFlag> stack = new Stack<>();
        stack.push(new ColorFlag(root, Constant.WHITE));
        while (!stack.isEmpty()) {
            ColorFlag colorFlag = stack.pop();
            BinaryTreeNode node = colorFlag.node;
            String color = colorFlag.color;
            if (color.equals(Constant.WHITE)) {
                stack.push(new ColorFlag(node, Constant.BLACK));
                if (node.right != null) {
                    stack.push(new ColorFlag(node.right, Constant.WHITE));
                }
                if (node.left != null) {
                    stack.push(new ColorFlag(node.left, Constant.WHITE));
                }
            }else {
                nums.add(node.val);
            }
        }
        return nums;
    }

    /**
     * 层序遍历
     */
    public List<List<Integer>> levelOrder() {
        List<List<Integer>> nums = new ArrayList<>();
        if (root == null) {
            return nums;
        }
        Queue<BinaryTreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            List<Integer> tmp = new ArrayList<>();
            int n = queue.size();
            for (int i = 0;i < n;i++) {
                BinaryTreeNode node = queue.poll();
                assert node != null;
                tmp.add(node.val);
                if (node.left != null) {
                    queue.offer(node.left);
                }
                if (node.right != null) {
                    queue.offer(node.right);
                }
            }
            nums.add(tmp);
        }
        return nums;
    }

    /**
     * Z型遍历
     */
    public List<List<Integer>> zOrder() {
        List<List<Integer>> nums = new ArrayList<>();
        if (root == null) {
            return nums;
        }
        Stack<BinaryTreeNode> positiveSequence = new Stack<>();
        Stack<BinaryTreeNode> backwardSequence = new Stack<>();
        positiveSequence.push(root);
        boolean isPositive = true;
        while (!positiveSequence.isEmpty() || !backwardSequence.isEmpty()) {
            List<Integer> tmp = new ArrayList<>();
            int n = (positiveSequence.size() != 0) ? positiveSequence.size(): backwardSequence.size();
            if (isPositive) {
                for (int i = 0;i < n;i++) {
                    BinaryTreeNode node = positiveSequence.pop();
                    tmp.add(node.val);
                    if (node.left != null) {
                        backwardSequence.push(node.left);
                    }
                    if (node.right != null) {
                        backwardSequence.push(node.right);
                    }
                }
            }else {
                for (int i = 0;i < n;i++) {
                    BinaryTreeNode node = backwardSequence.pop();
                    tmp.add(node.val);
                    if (node.right != null) {
                        positiveSequence.push(node.right);
                    }
                    if (node.left != null) {
                        positiveSequence.push(node.left);
                    }
                }
            }
            nums.add(tmp);
            isPositive = !isPositive;
        }
        return nums;
    }

    /**
     * 打印二叉树
     */
    public void printBinaryTreeNode(boolean isOrdinary) {
        List<List<Integer>> nums = (isOrdinary) ? levelOrder(): zOrder();
        for (List<Integer> num : nums) {
            System.out.println(num.toString());
        }
    }

    public static void main(String[] args) {
        BinaryTree binaryTree = UniversalMethod.createBinaryTree();

        System.out.println("----------------先序遍历（递归）----------------");
        List<Integer> preorderByRecurrence = binaryTree.preorderByRecurrence(new ArrayList<>(), binaryTree.root);
        System.out.println(preorderByRecurrence.toString());
        System.out.println("\n");

        System.out.println("----------------中序遍历（递归）----------------");
        List<Integer> inorderByRecurrence = binaryTree.inorderByRecurrence(new ArrayList<>(), binaryTree.root);
        System.out.println(inorderByRecurrence.toString());
        System.out.println("\n");

        System.out.println("----------------后序遍历（递归）----------------");
        List<Integer> postorderByRecurrence = binaryTree.postorderByRecurrence(new ArrayList<>(), binaryTree.root);
        System.out.println(postorderByRecurrence.toString());
        System.out.println("\n");

        System.out.println("----------------先序遍历（迭代）----------------");
        List<Integer> preorderByIteration = binaryTree.preorderByIteration();
        System.out.println(preorderByIteration.toString());
        System.out.println("\n");

        System.out.println("----------------中序遍历（迭代）----------------");
        List<Integer> inorderByIteration = binaryTree.inorderByIteration();
        System.out.println(inorderByIteration.toString());
        System.out.println("\n");

        System.out.println("----------------后序遍历（迭代）----------------");
        List<Integer> postorderByIteration = binaryTree.postorderByIteration();
        System.out.println(postorderByIteration.toString());
        System.out.println("\n");

        System.out.println("----------------层次遍历----------------");
        binaryTree.printBinaryTreeNode(true);
        System.out.println("\n");

        System.out.println("----------------Z型遍历----------------");
        binaryTree.printBinaryTreeNode(false);
        System.out.println("\n");
    }
}